import json
import socket
import os
import threading
import time
import io

CHUNK_SIZE=2048*1024#2Mb
TRACKER_IP='192.168.56.106'
PORT=6666
SOCKET_BUFFER_SIZE=1024
LOCAL_IP='192.168.56.1'
MEDIA_TYPE='.mp4'

peer_ip_list=[]

def send_to_tracker():
    file_list=os.listdir()
    file_info=[]
    for file in file_list:
        if MEDIA_TYPE in file:
            size=os.path.getsize(file)
            file_info.append({'file_name':file,'file_size':size})
    register={'request_type':1,'file_list':file_info}
    register=json.dumps(register).encode()
    sock=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
    sock.connect((TRACKER_IP,PORT))
    print('Send files information to the server...')

    for i in range(len(register)//SOCKET_BUFFER_SIZE+1):
        start=i*SOCKET_BUFFER_SIZE
        end=min(len(register),start+SOCKET_BUFFER_SIZE-1)
        sock.send(register[start:end])

    recv=sock.recv(SOCKET_BUFFER_SIZE)
    recv=json.loads(recv.decode())
    print(recv['status_code'])
    sock.close()

class recv_file(threading.Thread):
    def __init__(self,file_name,num,ip,start_chunk,end_chunk,size,chunk_num):
        threading.Thread.__init__(self)
        self.file_name=file_name
        self.num=num
        self.ip=ip
        self.start_chunk=start_chunk
        self.end_chunk=end_chunk
        self.chunk_num=chunk_num
        self.size=size

    def run(self):
        file=open(self.file_name+'.'+str(self.num),'wb')
        sock=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
        sock.connect((self.ip,PORT))
        request={'request_type':1,'file_name':self.file_name,'start_chunk':self.start_chunk,'end_chunk':self.end_chunk}
        request=json.dumps(request).encode()
        sock.send(request)
        
        t=time.time()
        while True:
            data=sock.recv(SOCKET_BUFFER_SIZE)
            if len(data)>0:
                file.write(data)
            else:
                break
        t=time.time()-t
        print('Segment {} downloaded, it cost {} second(s)'.format(self.num,t))
        file.close()
        sock.close()

        sum=0
        for i in range(self.chunk_num):
            chunk_file_name=self.file_name+'.'+str(i)
            sum+=os.path.getsize(chunk_file_name)
        if sum==self.size:
            print('All segments downloaded, now merge them')
            file=open(self.file_name,'wb')
            for i in range(self.chunk_num):
                chunk_file_name=self.file_name+'.'+str(i)
                chunk_size=os.path.getsize(chunk_file_name)
                chunk=open(chunk_file_name,'rb')
                while chunk.tell()<chunk_size:
                    data=chunk.read(CHUNK_SIZE)
                    file.write(data)
                chunk.close()
                os.remove(chunk_file_name)
            file.close()
            print('Merger complete')


class request_file(threading.Thread):
    def __init__(self,file_name):
        threading.Thread.__init__(self)
        self.file_name=file_name

    def run(self):
        sock=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
        sock.connect((TRACKER_IP,PORT))
        request={'request_type':0,'file_name':self.file_name}
        request=json.dumps(request).encode()
        sock.send(request)
        response=sock.recv(SOCKET_BUFFER_SIZE)
        sock.close()
        response=json.loads(response.decode())
        if not response['status_code']==200:
            print('No one has this file!')
            return

        size=response['file_size']
        IP_list=response['IP_list']

        global peer_ip_list
        for ip in IP_list:
            if not ip==LOCAL_IP:
                peer_ip_list.append(ip)

        total_chunk_num=size//CHUNK_SIZE+1
        peer_number=len(peer_ip_list)
        len_per_peer=(peer_number-1+total_chunk_num)//peer_number
        print('There are {} segments need to download...'.format(peer_number))
        for i in range(peer_number):
            start_chunk=i*len_per_peer
            end_chunk=min(total_chunk_num-1,start_chunk+len_per_peer-1)
            thread=recv_file(self.file_name,i,peer_ip_list[i],start_chunk,end_chunk,size,peer_number)
            thread.start()

class wait_input(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)

    def run(self):
        while True:
            message=input('please input here:\n')
            sock=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
            print('Send to...:\n')
            
            peer_num=len(peer_ip_list)
            for i in range(peer_num):
                print('{}. {}'.format(i,peer_ip_list[i]))
            print('{}. all peers'.format(peer_num))
            op=int(input())

            if op==peer_num:
                sock.connect((TRACKER_IP,PORT))
            else:
                sock.connect((peer_ip_list[op],PORT))
            request={'request_type':2,'message':message}
            request=json.dumps(request).encode()
            sock.send(request)
            sock.close()

class my_daemon(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
    def run(self):
        sock=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
        sock.bind(('',PORT))
        sock.listen()
    
        while True:
            conn,addr=sock.accept()
            recv=conn.recv(SOCKET_BUFFER_SIZE)
            recv=json.loads(recv.decode())

            if addr[0]==TRACKER_IP and not recv['source']==LOCAL_IP:
                print('\nMessage from {}:\n{}\n'.format(recv['source'],recv['message']))
            
            elif recv['request_type']==2:
                print('\nMessage from {}:\n{}\n'.format(addr[0],recv['message']))


file_name='test.mp4'
send_to_tracker()
request_file(file_name).start()
wait_input().start()
my_daemon().start()